import pandas as pd
import pickle as pkl
import numpy as np
import bidict
import random
import os
random.seed(1234)

def daytoweek(x):
    if x > "2020-09-22": return 105
    elif x > "2020-09-15": return 104
    elif x > "2020-09-08": return 103
    elif x > "2020-09-15": return 102
    elif x > "2020-09-01": return 101
    elif x > "2020-08-25": return 100
    elif x > "2020-08-18": return 99
    elif x > "2020-08-11": return 98
    elif x > "2020-08-04": return 97
    elif x > "2020-07-28": return 96
    elif x > "2020-07-21": return 95
    else: return 0

df = pd.read_csv(".\\input\\processed\\transactions_translated.csv", sep=',',index_col=0, usecols=[0, 1, 2, 3])
df.t_dat = df.t_dat.map(daytoweek)

item_feat = {}
user_feat = {}
with open('.\\input\\processed\\article_sparse.pkl', 'rb') as f:
    item_feat['artid'] = pkl.load(f)
    item_feat['pcode'] = pkl.load(f)
    artid_bidict = pkl.load(f)
    pcode_bidict = pkl.load(f)
    artid_cnt, pcode_cnt = pkl.load(f)

with open('.\\input\\processed\\customer_feat.pkl', 'rb') as f:
    user_feat['customer_id'] = pkl.load(f)
    user_feat['FN'] = pkl.load(f)
    user_feat['Active'] = pkl.load(f)
    user_feat['club_member_status'] = pkl.load(f)
    user_feat['fashion_news_frequency'] = pkl.load(f)
    user_feat['age'] = pkl.load(f)
    customer_cnt = pkl.load(f)
    customer_bidict = pkl.load(f)

with open('.\\input\\processed\\cf_matrices\\coclicks.pkl', 'rb') as f:
    cc = pkl.load(f)

train_set, valid_set, test_set = [], [], []
cnt = 0
for prev_week in range(95, 104, 1):
    fil = df.loc[df['t_dat'] == prev_week]
    fil_curr = df.loc[df['t_dat'] == prev_week + 1]
    for _, cf0 in fil.groupby('customer_id'):
        user_id = cf0.customer_id.iloc[0]
        cf1 = pd.unique(cf0.article_id).tolist()
        cf2 = [x for x in cf1]
        for x in cf1:
            for y in cc[prev_week][x]:
                cf2.append(y[1])
        cf2 = set(cf2)
        cf3 = set(fil_curr.loc[fil_curr['customer_id'] == user_id].article_id.tolist())
        pos_artid = cf2.intersection(cf3)
        neg_artid = cf2.difference(cf3)
        user_part = tuple(x[user_id] for x in user_feat.values())
        for artid in pos_artid:
            item_part = tuple(x[artid-1] for x in item_feat.values())
            train_set.append(user_part + item_part + (1, ))
            cnt += 1
        for artid in neg_artid:
            item_part = tuple(x[artid-1] for x in item_feat.values())
            train_set.append(user_part + item_part + (0, ))
    print(prev_week, "Done")
random.shuffle(train_set)
piv = len(train_set) // 10
train_set, valid_set = train_set[:-piv], train_set[-piv:]
print(cnt, len(train_set), len(valid_set))

prev_week = 104
fil = df.loc[df['t_dat'] == prev_week]
fil_hist = df.loc[df['t_dat'] > 0]
for _, cf0 in fil.groupby('customer_id'):
    user_id = cf0.customer_id.iloc[0]
    cf1 = fil[fil['customer_id'] == user_id].article_id.tolist()
    cf2 = fil_hist[fil_hist['customer_id'] == user_id].article_id.tolist()
    for x in cf1:
        for y in cc[prev_week][x]:
            cf2.append(y[1])
    cf2 = set(cf2)
    user_part = tuple(x[user_id] for x in user_feat.values())
    for artid in cf2:
        item_part = tuple(x[artid-1] for x in item_feat.values())
        test_set.append(user_part + item_part + (1, ))

print(len(test_set))

with open('.\\input\\processed\\train_valid.pkl', 'wb') as f:
    pkl.dump(train_set, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(valid_set, f, pkl.HIGHEST_PROTOCOL)

with open('.\\input\\processed\\test.pkl', 'wb') as f:
    pkl.dump(test_set, f, pkl.HIGHEST_PROTOCOL)
